import taichi as ti
import numpy as np

ti.init()
n=4
m=7
# taichi tensors
val=ti.var(ti.i32,shape=(n,m))
vec=ti.Vector(3,dt=ti.i32,shape=(n,m))
mat=ti.Matrix(3,4,dt=ti.i32,shape=(n,m))

# scalar
arr=np.ones(shape=(n,m),dtype=np.int32)
val.from_numpy(arr)
arr=val.to_numpy()

# vector
arr=np.ones(shape=(n,m,3),dtype=np.int32)
vec.from_numpy(arr)
arr=np.ones(shape=(n,m,3,1),dtype=np.int32)
vec.from_numpy(arr)
arr = vec.to_numpy()
assert arr.shape == (n,m,3)

# matrix
arr=np.ones(shape=(n,m,3,4),dtype=np.int32)
mat.from_numpy(arr)

arr=mat.to_numpy()
assert arr.shape==(n,m,3,4)